package per.jdz.demo.utils;

import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.Resource;
import java.util.Collection;
import java.util.function.BiFunction;

@Slf4j
@Component
public class MybatisBatchUtil {
    
    /**
    * 每次处理1000条
    */
    private static final int BATCH_SIZE = 1000;
    
    @Resource
    private SqlSessionFactory sqlSessionFactory;
    
    /**
    * 批量处理修改或者插入
    *
    * @param data     需要被处理的数据
    * @param mapperClass  Mybatis的Mapper类
    * @param function 自定义处理逻辑
    * @return int 影响的总行数
    */
    @Transactional(rollbackFor = Exception.class)
    public <T, U, R> int batchUpdateOrInsert(Collection<T> data, Class<U> mapperClass, BiFunction<T, U, R> function) {
        int count = 1;
        try (SqlSession batchSqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH)) {
            U mapper = batchSqlSession.getMapper(mapperClass);
            int size = data.size();
            for (T element : data) {
                if (function == null) {
                    log.error("没匹配到公共数据");
                    return 0;
                }
                function.apply(element, mapper);
                if ((count % BATCH_SIZE == 0) || count == size) {
                    batchSqlSession.flushStatements();
                }
                count++;
            }
        } catch (Exception e) {
            log.error("批量插入发生异常：{}", e.getMessage(), e);
            throw new RuntimeException(e.getMessage());
        }
        return count - 1;
    }
}
